{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d4ffa253-e427-4663-bc0b-9b1fcce508e6",
   "metadata": {
    "id": "d4ffa253-e427-4663-bc0b-9b1fcce508e6"
   },
   "source": [
    "# 路由\n",
    "\n",
    "## 环境\n",
    "\n",
    "`(1) 包`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc509da-52b2-49fc-bc45-e5fd75ff5fed",
   "metadata": {
    "id": "9bc509da-52b2-49fc-bc45-e5fd75ff5fed",
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install -q langchain_community tiktoken langchain-openai langchainhub  langchain\n",
    "! pip install -q chromadb==0.4.15\n",
    "! pip install -q beautifulsoup4\n",
    "! pip install -q langchain-nomic\n",
    "! pip install -q --upgrade httpx httpx-sse PyJWT\n",
    "! pip install -q --upgrade --quiet  dashscope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "IyaZnM90Nn_1",
   "metadata": {
    "id": "IyaZnM90Nn_1"
   },
   "outputs": [],
   "source": [
    "from google.colab import userdata\n",
    "DASHSCOPE_API_KEY=userdata.get('DASHSCOPE_API_KEY')\n",
    "DEEPSEEK_API_KEY=userdata.get('DEEPSEEK_API_KEY')\n",
    "LANGCHAIN_API_KEY=userdata.get('LANGCHAIN_API_KEY')\n",
    "OPENAI_API_KEY=userdata.get('OPENAI_API_KEY')\n",
    "\n",
    "from langchain_community.embeddings import DashScopeEmbeddings\n",
    "import os\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = f\"RAG_路由\"\n",
    "os.environ['LANGCHAIN_TRACING_V2'] = 'true'\n",
    "os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'\n",
    "os.environ['LANGCHAIN_API_KEY'] = LANGCHAIN_API_KEY\n",
    "os.environ['USER_AGENT'] = 'myagent'\n",
    "os.environ[\"DASHSCOPE_API_KEY\"] = DASHSCOPE_API_KEY\n",
    "os.environ[\"DEEPSEEK_API_KEY\"] = DEEPSEEK_API_KEY\n",
    "os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8a28de58-e6f9-4e29-8fd9-89cb47096a26",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:32:12.028590Z",
     "iopub.status.busy": "2024-10-28T03:32:12.028321Z",
     "iopub.status.idle": "2024-10-28T03:32:13.664059Z",
     "shell.execute_reply": "2024-10-28T03:32:13.663657Z",
     "shell.execute_reply.started": "2024-10-28T03:32:12.028575Z"
    },
    "id": "8a28de58-e6f9-4e29-8fd9-89cb47096a26",
    "outputId": "242927be-46b2-4059-cafc-6da63a04d414",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "你好！有什么我可以帮助你的吗？\n"
     ]
    }
   ],
   "source": [
    "# from langchain_community.chat_models import ChatZhipuAI\n",
    "# from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "\n",
    "# llm = ChatZhipuAI(model=\"glm-4-plus\",temperature=0.5)\n",
    "# messages = [\n",
    "#     AIMessage(content=\"Hi.\"),\n",
    "#     SystemMessage(content=\"Your role is a poet.\"),\n",
    "#     HumanMessage(content=\"Write a short poem about AI in four lines.\"),\n",
    "# ]\n",
    "\n",
    "# response = llm.invoke(messages)\n",
    "# print(response.content)  # Displays the AI-generated poem\n",
    "\n",
    "# from langchain_community.chat_models.tongyi import ChatTongyi\n",
    "# from langchain_core.messages import HumanMessage\n",
    "\n",
    "# llm = ChatTongyi(\n",
    "#     streaming=True,\n",
    "# )\n",
    "# res = llm.stream([HumanMessage(content=\"hi\")], streaming=True)\n",
    "# for r in res:\n",
    "#     print(\"chat resp:\", r)\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "llm=ChatOpenAI(model=\"gpt-4o\")\n",
    "# llm = ChatOpenAI(\n",
    "#     model='deepseek-reasoner',\n",
    "#     openai_api_key=DEEPSEEK_API_KEY,\n",
    "#     openai_api_base='https://api.deepseek.com'\n",
    "# )\n",
    "\n",
    "response = llm.invoke(\"你好!\")\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45fd9558-c53f-4a6c-80f0-c31b3bfd55de",
   "metadata": {
    "id": "45fd9558-c53f-4a6c-80f0-c31b3bfd55de"
   },
   "source": [
    "## 第 10 部分：逻辑和语义路由\n",
    "\n",
    "建立几个公司部门，规定他们的职责\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "04c2cf60-d636-4992-a021-1236f7688999",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:22:06.484037Z",
     "iopub.status.busy": "2024-10-28T03:22:06.483768Z",
     "iopub.status.idle": "2024-10-28T03:22:06.532248Z",
     "shell.execute_reply": "2024-10-28T03:22:06.531861Z",
     "shell.execute_reply.started": "2024-10-28T03:22:06.484020Z"
    },
    "id": "04c2cf60-d636-4992-a021-1236f7688999",
    "outputId": "0ee3e1ed-2616-4f2c-eee4-b44855159337",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py:3553: LangChainDeprecationWarning: As of langchain-core 0.3.0, LangChain uses pydantic v2 internally. The langchain_core.pydantic_v1 module was a compatibility shim for pydantic v1, and should no longer be used. Please update the code to import from Pydantic directly.\n",
      "\n",
      "For example, replace imports like: `from langchain_core.pydantic_v1 import BaseModel`\n",
      "with: `from pydantic import BaseModel`\n",
      "or the v1 compatibility namespace if you are working in a code base that has not been fully upgraded to pydantic 2 yet. \tfrom pydantic.v1 import BaseModel\n",
      "\n",
      "  exec(code_obj, self.user_global_ns, self.user_ns)\n",
      "/usr/local/lib/python3.11/dist-packages/langchain_openai/chat_models/base.py:1362: UserWarning: Received a Pydantic BaseModel V1 schema. This is not supported by method=\"json_schema\". Please use method=\"function_calling\" or specify schema via JSON Schema or Pydantic V2 BaseModel. Overriding to method=\"function_calling\".\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# 数据模型部门\n",
    "class RouteQuery(BaseModel):\n",
    "    \"\"\"将用户查询路由到最相关的数据源。\"\"\"\n",
    "\n",
    "    datasource: Literal[\"python_docs\", \"js_docs\", \"golang_docs\"] = Field(\n",
    "        ...,\n",
    "        description=\"给出一个用户问题，选择哪个数据源与回答他们的问题最相关\",\n",
    "    )\n",
    "\n",
    "# 带函数调用的LLM\n",
    "structured_llm = llm.with_structured_output(RouteQuery)\n",
    "\n",
    "# Prompt\n",
    "system = \"\"\"您是将用户问题路由到适当数据源的专家。\n",
    "\n",
    "根据问题所指的编程语言，将其路由到相关数据源。\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# 定义路由器\n",
    "router = prompt | structured_llm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cb5aaa3-7826-471c-9203-390009cbb4d6",
   "metadata": {
    "id": "1cb5aaa3-7826-471c-9203-390009cbb4d6"
   },
   "source": [
    "客户来了，问了一个问题应该由某个部门去负责"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cfc6febc-93df-49b4-9920-c93589ba021e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-28T03:23:23.168788Z",
     "iopub.status.busy": "2024-10-28T03:23:23.168526Z",
     "iopub.status.idle": "2024-10-28T03:23:25.154383Z",
     "shell.execute_reply": "2024-10-28T03:23:25.153942Z",
     "shell.execute_reply.started": "2024-10-28T03:23:23.168772Z"
    },
    "id": "cfc6febc-93df-49b4-9920-c93589ba021e",
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"\"\"为什么下面的代码不起作用:\n",
    "\n",
    "import random\n",
    "import string\n",
    "\n",
    "def generate_password(length):\n",
    "    characters = string.ascii_letters + string.digits + string.punctuation\n",
    "    password = ''.join(random.choice(characters) for i in range(length))\n",
    "    return password\n",
    "\n",
    "password = generate_password(12)\n",
    "print(\"Generated password:\", password)\n",
    "\"\"\"\n",
    "\n",
    "result = router.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "277536df-0904-4d99-92bb-652621afbdec",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:23:27.894415Z",
     "iopub.status.busy": "2024-10-28T03:23:27.894156Z",
     "iopub.status.idle": "2024-10-28T03:23:27.898962Z",
     "shell.execute_reply": "2024-10-28T03:23:27.898570Z",
     "shell.execute_reply.started": "2024-10-28T03:23:27.894400Z"
    },
    "id": "277536df-0904-4d99-92bb-652621afbdec",
    "outputId": "9f493ab2-7ee7-414e-d9ac-2ab0f29a858b",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RouteQuery(datasource='python_docs')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "636a43ae-50f3-43a1-a1b7-93266ea13bcd",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:23:30.887504Z",
     "iopub.status.busy": "2024-10-28T03:23:30.887233Z",
     "iopub.status.idle": "2024-10-28T03:23:30.890450Z",
     "shell.execute_reply": "2024-10-28T03:23:30.890063Z",
     "shell.execute_reply.started": "2024-10-28T03:23:30.887487Z"
    },
    "id": "636a43ae-50f3-43a1-a1b7-93266ea13bcd",
    "outputId": "8bbea82e-e24b-49b3-e8ea-4af311cb612b",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'python_docs'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.datasource"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba569c12-ba1d-4486-a5be-bc8c31a8448c",
   "metadata": {
    "id": "ba569c12-ba1d-4486-a5be-bc8c31a8448c"
   },
   "source": [
    "一旦我们有了这个，我们就可以很轻松的转到各个线条"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "01f15722-35c6-4456-ad1b-06463233db25",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-28T03:23:34.975228Z",
     "iopub.status.busy": "2024-10-28T03:23:34.974969Z",
     "iopub.status.idle": "2024-10-28T03:23:34.978810Z",
     "shell.execute_reply": "2024-10-28T03:23:34.978447Z",
     "shell.execute_reply.started": "2024-10-28T03:23:34.975213Z"
    },
    "id": "01f15722-35c6-4456-ad1b-06463233db25",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def choose_route(result):\n",
    "    if \"python_docs\" in result.datasource.lower():\n",
    "        ### 逻辑在这里\n",
    "        return \"chain for python_docs\"\n",
    "    elif \"js_docs\" in result.datasource.lower():\n",
    "        ### 逻辑在这里\n",
    "        return \"chain for js_docs\"\n",
    "    else:\n",
    "        ### 逻辑在这里\n",
    "        return \"golang_docs\"\n",
    "\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "\n",
    "full_chain = router | RunnableLambda(choose_route)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6af07b77-0537-4635-87ec-ad8f59d34e9b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:23:37.780043Z",
     "iopub.status.busy": "2024-10-28T03:23:37.779782Z",
     "iopub.status.idle": "2024-10-28T03:23:39.492171Z",
     "shell.execute_reply": "2024-10-28T03:23:39.491760Z",
     "shell.execute_reply.started": "2024-10-28T03:23:37.780029Z"
    },
    "id": "6af07b77-0537-4635-87ec-ad8f59d34e9b",
    "outputId": "e38d01b7-7912-4393-b54e-608fd7f8b706",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'chain for python_docs'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_chain.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb1db843-95f4-4abd-8d4c-4a41f0910949",
   "metadata": {
    "id": "cb1db843-95f4-4abd-8d4c-4a41f0910949"
   },
   "source": [
    "### 语义路由\n",
    "不同的专家线条"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "53cbfa72-c35a-4d1d-aa6d-08a570ab2170",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-28T03:32:20.014341Z",
     "iopub.status.busy": "2024-10-28T03:32:20.014074Z",
     "iopub.status.idle": "2024-10-28T03:38:54.477687Z",
     "shell.execute_reply": "2024-10-28T03:38:54.476681Z",
     "shell.execute_reply.started": "2024-10-28T03:32:20.014327Z"
    },
    "id": "53cbfa72-c35a-4d1d-aa6d-08a570ab2170",
    "outputId": "85249f4a-4b89-4bd9-9f12-d2ed768deaff",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "呼叫物理专家\n",
      "黑洞是宇宙中一种非常神秘而强大的天体，它具有极大的引力，甚至连光线也无法逃脱其吸引。黑洞形成于星体坍缩的过程中，当一颗恒星燃烧全部的燃料后，会发生坍缩并形成黑洞。黑洞内部具有一个奇点，即一个极为奇特的空间点，其引力无限大，时间和空间也会产生扭曲。目前科学家对黑洞的研究还在探索中，但我们可以确定的是黑洞是宇宙中非常奇妙的存在。\n"
     ]
    }
   ],
   "source": [
    "from langchain.utils.math import cosine_similarity\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "\n",
    "# 两个 prompts\n",
    "physics_template = \"\"\"你是一位非常聪明的物理学教授。\\\n",
    "你擅长以简洁易懂的方式回答有关物理的问题。\\\n",
    "当你不知道问题的答案时，你会承认你不知道。\n",
    "\n",
    "这里有一个问题：\n",
    "{query}\"\"\"\n",
    "\n",
    "math_template = \"\"\"你是一位非常优秀的数学家。你非常擅长回答数学问题。\\\n",
    "你之所以如此优秀，是因为你能将难题分解成各个组成部分，\\\n",
    "回答各个组成部分的问题，然后将它们组合起来回答更广泛的问题。\n",
    "\n",
    "这里有一个问题：\n",
    "{query}\"\"\"\n",
    "\n",
    "# 嵌入\n",
    "embeddings = OpenAIEmbeddings()\n",
    "prompt_templates = [physics_template, math_template]\n",
    "prompt_embeddings = embeddings.embed_documents(prompt_templates)\n",
    "\n",
    "# 将问题转至提示\n",
    "def prompt_router(input):\n",
    "    # 嵌入问题\n",
    "    query_embedding = embeddings.embed_query(input[\"query\"])\n",
    "    # 计算相似度\n",
    "    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]\n",
    "    most_similar = prompt_templates[similarity.argmax()]\n",
    "    # 选择的提示\n",
    "    print(\"呼叫数学专家\" if most_similar == math_template else \"呼叫物理专家\")\n",
    "    return PromptTemplate.from_template(most_similar)\n",
    "\n",
    "\n",
    "chain = (\n",
    "    {\"query\": RunnablePassthrough()}\n",
    "    | RunnableLambda(prompt_router)\n",
    "    | ChatOpenAI()\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "print(chain.invoke(\"什么是黑洞？\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93d3516-3f77-4548-b55f-5db4d7c9fbbb",
   "metadata": {
    "id": "e93d3516-3f77-4548-b55f-5db4d7c9fbbb"
   },
   "source": [
    "## 结构化搜索\n",
    "\n",
    "\n",
    "许多矢量存储包含元数据字段。\n",
    "\n",
    "这使得可以根据元数据过滤特定块。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "KlB_8FxgTfOj",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KlB_8FxgTfOj",
    "outputId": "ef74bbee-1f2f-4a53-b653-e470eef20208"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/111.0 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r",
      "\u001b[2K   \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━\u001b[0m \u001b[32m102.4/111.0 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m111.0/111.0 kB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "%pip install --upgrade --quiet  lark langchain-chroma"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53767f87-f38d-4f16-9b3b-b83ce4657f03",
   "metadata": {
    "id": "53767f87-f38d-4f16-9b3b-b83ce4657f03"
   },
   "source": [
    "为了演示目的，我们将使用 Chroma 矢量存储。我们创建了一组包含电影摘要的小型演示文档。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7731745b-accc-4cf1-8291-e12d1aa46361",
   "metadata": {
    "id": "7731745b-accc-4cf1-8291-e12d1aa46361"
   },
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_core.documents import Document\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "docs = [\n",
    "    Document(\n",
    "        page_content=\"一群科学家带回恐龙，混乱随之而来\",\n",
    "        metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": \"science fiction\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"莱昂纳多·迪卡普里奥迷失在梦中梦中梦中梦中...\",\n",
    "        metadata={\"year\": 2010, \"director\": \"Christopher Nolan\", \"rating\": 8.2},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"一个心理学家/侦探迷失在梦中梦中梦中，盗梦空间重用了这个想法\",\n",
    "        metadata={\"year\": 2006, \"director\": \"Satoshi Kon\", \"rating\": 8.6},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"一群正常身材的女性极其健康，一些男性对她们心生向往\",\n",
    "        metadata={\"year\": 2019, \"director\": \"Greta Gerwig\", \"rating\": 8.3},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"玩具活了过来，并乐在其中\",\n",
    "        metadata={\"year\": 1995, \"genre\": \"animated\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"三个人走入区域，三个人走出区域\",\n",
    "        metadata={\n",
    "            \"year\": 1979,\n",
    "            \"director\": \"Andrei Tarkovsky\",\n",
    "            \"genre\": \"thriller\",\n",
    "            \"rating\": 9.9,\n",
    "        },\n",
    "    ),\n",
    "]\n",
    "vectorstore = Chroma.from_documents(docs, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6054bc98-0ae1-45e6-8f06-22c65ec47180",
   "metadata": {
    "id": "6054bc98-0ae1-45e6-8f06-22c65ec47180"
   },
   "source": [
    "## 创建我们的自查询检索器\n",
    "现在我们可以实例化我们的检索器。为此，我们需要预先提供一些关于我们的文档支持的元数据字段的信息以及文档内容的简短描述。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f699d9e7-468e-4574-bdba-f4be4a5779de",
   "metadata": {
    "id": "f699d9e7-468e-4574-bdba-f4be4a5779de"
   },
   "outputs": [],
   "source": [
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"genre\",\n",
    "        description=\"电影的类型。类型之一 ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"year\",\n",
    "        description=\"电影上映的年份\",\n",
    "        type=\"integer\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"director\",\n",
    "        description=\"电影导演的名字\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"rating\", description=\"电影的1-10评分\", type=\"float\"\n",
    "    ),\n",
    "]\n",
    "document_content_description = \"电影的简要总结\"\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectorstore,\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-_VQafSwY1V6",
   "metadata": {
    "id": "-_VQafSwY1V6"
   },
   "source": [
    "## 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1b776858-a589-4fe5-a8a3-19530706075d",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1b776858-a589-4fe5-a8a3-19530706075d",
    "outputId": "af89f3e6-40fe-4e56-d6ea-d13feed0d2a1"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(id='6e8686e6-882a-407a-b090-2ae856c595d1', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006}, page_content='一个心理学家/侦探迷失在梦中梦中梦中，盗梦空间重用了这个想法'),\n",
       " Document(id='1b4048ae-93ee-4802-9e98-02b1afd11a67', metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979}, page_content='三个人走入区域，三个人走出区域')]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.invoke(\"我想看一部评分高于 8.5 的电影\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "65bfad12-1985-433e-a980-eb8c9da53f72",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "65bfad12-1985-433e-a980-eb8c9da53f72",
    "outputId": "de0fec65-4454-4d22-a72d-cc473d9f4432"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(id='a8c8f2f8-bd01-4653-9a3e-0ffe716d938c', metadata={'director': 'Greta Gerwig', 'rating': 8.3, 'year': 2019}, page_content='一群正常身材的女性极其健康，一些男性对她们心生向往')]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.invoke(\"Greta Gerwig导演过任何关于女性的电影吗\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "99643372-01cc-49cc-a507-bebbed096247",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "99643372-01cc-49cc-a507-bebbed096247",
    "outputId": "b19ca154-2eda-4f67-ff15-9ee1f1364a62"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(id='caba05c5-4946-475f-8902-118defe78839', metadata={'genre': 'animated', 'year': 1995}, page_content='玩具活了过来，并乐在其中')]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever.invoke(\n",
    "    \"1990 年之后、2005 年之前的电影都与玩具有关，最好是动画片\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
